package com.codemes.happylist.lucene.knn;


import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.util.Arrays;
import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.VectorUtil;

/**
 * Looks up each tokens in a dictionary, and sums the token vectors. Unrecognized tokens are
 * ignored. The resulting vector is normalized to unit length.
 */
public final class KnnVectorDictFilter extends TokenFilter {

    private final KnnVectorDict dict;
    private final CharTermAttribute termAtt = addAttribute(CharTermAttribute.class);
    private final float[] scratchFloats;
    private final float[] result;
    private final byte[] scratchBytes;
    private final FloatBuffer scratchBuffer;

    /**
     * sole constructor
     *
     * @param input the input token stream to filter.
     * @param dict a token to vector dictionary, used to look up the token vectors.
     */
    public KnnVectorDictFilter(TokenStream input, KnnVectorDict dict) {
        super(input);
        this.dict = dict;
        result = new float[dict.getDimension()];
        scratchBytes = new byte[dict.getDimension() * Float.BYTES];
        scratchBuffer = ByteBuffer.wrap(scratchBytes).order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer();
        scratchFloats = new float[dict.getDimension()];
    }

    @Override
    public boolean incrementToken() throws IOException {
        if (input.incrementToken() == false) {
            return false;
        }
        BytesRef term = new BytesRef(termAtt.toString());
        dict.get(term, scratchBytes);
        scratchBuffer.position(0);
        scratchBuffer.get(scratchFloats);
        VectorUtil.add(result, scratchFloats);
        return true;
    }

    @Override
    public void reset() throws IOException {
        super.reset();
        Arrays.fill(result, 0);
    }

    @Override
    public void end() throws IOException {
        super.end();
        VectorUtil.l2normalize(result, false);
    }

    /**
     * Get the vector computed from the input
     *
     * @return the resultant sum of the vectors of each term.
     */
    public float[] getResult() {
        return result;
    }
}